import torch
import torch.nn as nn

x1=0.0
x2 = 1.0
x = torch.tensor([
    [x1,x2],
    [x1,x2],
[x2,x1],
[x2,x1],
[x2,x1]
])
y = nn.functional.softmax(x,dim=1)
print(x.shape)
print(y)
# c = torch.multinomial(y,3)
# print(c)

func = torch.nn.LogSoftmax(dim=1)
y = func(x)
L = nn.NLLLoss()
z = torch.tensor([1,1,0,0,0])
loss = L(y,z)
print(loss)